import numpy as np
from sko.DE import DE
from pyted.frontend import Variable
from pyted import get_circuit, DataSweep, debug
import copy
from sko.tools import func_transformer
from torch import Tensor
from frontend.torch import TorchVariable
import sys
import time


class LSHADE():

    def __init__(self, fitness, lowwer, upper, pop_size, dim, epochs, tracer, eval):
        np.random.seed(int(time.time()))
        self.eval = np.array(eval)
        self.fitness = fitness
        self.tracer = tracer
        self.lowbound = np.array(lowwer)
        self.upbound = np.array(upper)
        self.pop_size = pop_size
        self.dim = dim
        self.population = np.random.uniform(
            self.lowbound, self.upbound, (self.pop_size, self.dim))
        self.fit = np.random.rand(self.pop_size)
        self.conv = np.random.rand(self.pop_size)
        self.best = self.population[0]
        self.Archive = np.empty_like(self.population)
        self.Archive[:] = np.nan
        self.fit_ea = np.random.rand(self.pop_size)
        self.conv_ea = np.random.rand(self.pop_size)
        self.time_up = 2*60*60 - 200
        self.start_time = time.time()

        self.uCR = 0.5
        self.uF = 0.5
        self.F = np.array([0.5 for _ in range(self.pop_size)])
        self.CR = np.array([0.5 for _ in range(self.pop_size)])
        self.c = 0.01
        self.SCR = []
        self.SF = []
        self.success_set = []
        self.fail_set = []
        self.H = 6
        self.MCR = np.array([0.5] * self.H)
        self.MF = np.array([0.5] * self.H)
        self.stopvalue = 0
        # ***************************************
        self.Epochs = epochs

    def initpop(self):
        # self.population = self.lowbound + \
        #     (self.upbound - self.lowbound) * \
        #     np.random.rand(self.pop_size, self.dim)
        self.population = self.lowbound + \
            (self.upbound - self.lowbound) * \
            np.random.rand(self.pop_size-1, self.dim)
        self.population = np.concatenate(
            (self.population, self.eval.reshape(1, -1)), axis=0)
        self.fit, self.conv = self.fitness(self.population)
        self.best = self.population[np.argmin(self.fit)]

    def mut(self):
        p = max(int(self.pop_size * 0.1), 1)
        # idx_set = np.argsort(self.fit)[:p]
        idx_set = np.lexsort((self.conv, self.fit))[:p]
        xpbest_idx = np.random.choice(idx_set, self.pop_size, replace=True)
        xpbest = self.population[xpbest_idx].reshape(
            -1, self.population.shape[1])
        a_indices = np.random.choice(
            np.delete(np.arange(self.pop_size), xpbest_idx), self.pop_size, replace=True)
        a = self.population[a_indices].reshape(-1, self.population.shape[1])
        combined_population = np.concatenate(
            (self.population, self.Archive), axis=0)
        combined_population = combined_population[~np.all(
            combined_population == 0, axis=1)]
        b_indices = np.random.choice(
            len(combined_population), self.pop_size, replace=True)
        b = combined_population[b_indices]
        F_broadcasted = np.tile(self.F[:, np.newaxis], (1, self.dim))
        v = self.population + F_broadcasted * \
            (xpbest - self.population) + F_broadcasted * (a - b)

        return v

    # def mut(self):
    #     r1 = np.random.choice(self.pop_size, size=len(
    #         self.population), replace=True)
    #     r2 = np.random.choice(self.pop_size, size=len(
    #         self.population), replace=True)
    #     r3 = np.random.choice(self.pop_size, size=len(
    #         self.population), replace=True)

    #     r2 = np.where(r2 == r1, (r1 + 1) % self.pop_size, r2)
    #     r3 = np.where(np.logical_or(r3 == r1, r3 == r2),
    #                   (r2 + 1) % self.pop_size, r3)
    #     conv_broadcasted = np.expand_dims(self.conv, axis=1)
    #     condition = conv_broadcasted[r1] == 0 & ~np.isnan(self.Archive[r2])
    #     p = np.where(condition, self.Archive[r2], self.population[r2])
    #     F_broadcasted = np.tile(self.F[:, np.newaxis], (1, self.dim))
    #     v = self.population[r1] - F_broadcasted * (p - self.population[r3])

    #     return v

    def cross(self, mut_chorm):
        j = np.random.randint(0, self.dim, size=self.pop_size)
        CR_broadcasted = np.tile(self.CR[:, np.newaxis], (1, self.dim))
        crossover_mask = (np.random.rand(self.pop_size, self.dim) < CR_broadcasted) | (
            np.arange(self.dim) == j[:, np.newaxis])
        cross_chorm = np.where(crossover_mask, mut_chorm, self.population)

        nan_mask = np.isnan(cross_chorm)
        out_of_bounds = (cross_chorm < self.lowbound) | (
            cross_chorm > self.upbound) | nan_mask
        new_values = np.random.uniform(
            self.lowbound, self.upbound, size=cross_chorm.shape)
        cross_chorm[out_of_bounds] = new_values[out_of_bounds]

        return cross_chorm

    def select(self, cross_chorm):
        # with SuppressOutput():
        temp, temp_v = self.fitness(cross_chorm)
        # temp_v = self.constraints(cross_chorm)

        condition_1 = (temp < self.fit) & (self.conv < temp_v)
        condition_2 = np.isnan(self.Archive).all(axis=1)
        condition_3 = (self.fit_ea < temp) & (temp_v < self.conv_ea)
        update_condition = condition_1 & (condition_2 | condition_3)
        self.Archive[update_condition] = np.where(
            condition_2[:, np.newaxis], cross_chorm, self.Archive)[update_condition]

        self.fit_ea[update_condition] = np.where(
            condition_2, temp, self.fit)[update_condition]
        self.conv_ea[update_condition] = np.where(
            condition_2, temp_v, self.conv)[update_condition]
        reset_condition = (temp < self.fit_ea) & (temp_v <= self.fit)
        self.Archive[reset_condition] = np.nan

        mask_1 = (self.conv == temp_v) & (self.fit >= temp)
        mask_2 = (self.conv > temp_v)
        mask = mask_1 | mask_2

        self.SCR.extend(self.CR[mask])
        self.SF.extend(self.F[mask])
        self.success_set.extend(np.where(mask)[0])
        self.fail_set.extend(self.fit[mask])

        self.population[mask] = cross_chorm[mask]
        self.fit[mask] = temp[mask]
        self.conv[mask] = temp_v[mask]

    def run(self):
        pre = np.min(self.fit)
        count = 0
        self.initpop()
        self.last_min_fitness_value = 0
        self.unimprove = 0

        # random.seed(current_time)

        for epoch in range(self.Epochs):
            # current_time = int(time.time())
            # np.random.seed(current_time)
            self.unimprove += 1
            self.SCR = []
            self.SF = []
            self.success_set = []
            self.fail_set = []
            mcr_indices = np.random.choice(len(self.MCR), self.pop_size)
            mf_indices = np.random.choice(len(self.MF), self.pop_size)
            self.CR = np.where(self.MCR[mcr_indices] == self.stopvalue, 0, np.random.normal(
                self.MCR[mcr_indices], 0.1))
            self.CR = np.clip(self.CR, 0, 1)
            self.F = self.MF[mf_indices] + \
                np.sqrt(0.1) * np.random.standard_cauchy(self.pop_size)
            self.F[self.F <= 0] = np.random.standard_cauchy(
                np.sum(self.F <= 0))
            self.F[self.F > 1] = 1
            mut_chorms = self.mut()
            cross_chorms = self.cross(mut_chorms)
            self.select(cross_chorms)

            # self.best = self.population[np.argmin(self.fit)]
            # min_conv = np.min(self.conv)
            # min_fit = np.min(self.fit[self.conv == min_conv])
            min_constraint_indices = np.where(
                self.conv == np.min(self.conv))[0]
            min_fitness_among_min_constraint = np.argmin(
                self.fit[min_constraint_indices])
            min_fitness_index = min_constraint_indices[min_fitness_among_min_constraint]
            self.best = self.population[min_fitness_index]
            min_fitness_value = self.fit[min_fitness_index]
            min_constraint_value = self.conv[min_fitness_index]
            self.tracer.collect_n_iters_fom(
                min_fitness_value,
                np.array([min_constraint_value]),
            )

            if self.last_min_fitness_value != min_fitness_value:
                self.unimprove = 0
            self.last_min_fitness_value = copy.deepcopy(min_fitness_value)
            # if min_constraint_value == 0 and self.unimprove >= 1:
            #     # print(min_constraint_value, min_fitness_value)
            #     break
            if min_constraint_value == 0:
                # print(min_constraint_value, min_fitness_value)
                break
            if (np.any(self.SCR) and np.any(self.SF)):
                d_f = np.abs(self.fail_set - self.fit[self.success_set])
                w = d_f / np.sum(d_f)
                self.uCR = np.sum(np.multiply(w, np.array(
                    self.SCR) ** 2)) / np.sum(np.multiply(w, np.array(self.SCR)))
                self.uF = np.sum(np.multiply(w, np.array(self.SF) ** 2)) / \
                    np.sum(np.multiply(w, np.array(self.SF)))
                self.MCR[epoch % self.H] = self.stopvalue if (
                    self.MF[epoch % self.H] == self.stopvalue or max(self.SCR) == 0) else self.uCR
                self.MF[epoch % self.H] = self.uF
            # ft, cv = self.fitness(self.best)
            # print(np.min(self.conv), end='\t')
            # print('epoch:', epoch, 'best:', np.min(self.fit),
            #       'mean:', np.mean(self.fit),
            #       'worst:', np.max(self.fit),
            #       'std:', np.std(self.fit),
            #       #   'NFE:', self.NFE,
            #       'pop_size', self.pop_size,
            #       "min_fit:", min_fitness_value,
            #       "min_conv:", min_constraint_value,
            #       'time:', time.time()-self.start_time)
            if time.time() - self.start_time > self.time_up:
                break
        # print(self.best)
        # print(type(self.best))
        # print(self.best.shape)
        # tmp = self.obj_func(self.best.reshape(1,6))
        # tmp = self.constraints(self.best.reshape(1,6))
        return self.best, self.tracer


class Test_Demo:
    def __init__(
        self,
        tester,
        problem,
        tracer=None,
        ref_point=None,
        pop_size=50,
        maxiters=1000,
    ):
        # set the parameters
        # =============================#
        # =============================#
        self.problem = problem
        self.tester = tester
        self.pop_size = pop_size
        self.maxiters = maxiters
        self.tracer = tracer(ref_point)
        Dim = len(get_circuit().variables)
        # bound = [(v._lb, v._ub) for _, v in get_circuit().variables.items()]
        # lb, ub = map(list, zip(*bound))
        bound = [(v._lb, v._ub, v._eval)
                 for _, v in get_circuit().variables.items()]
        lb, ub, eval = map(list, zip(*bound))

        def obj_func(x):
            # # WRITE CODE HERE within two '#' bar                                        #
            # # Converts the input x to list form and dim(x)=1                            #
            # #############################################################################

            # #############################################################################
            tester = self.update(x)
            prob = self.problem(tester)

            self._obj = prob.objective[0]
            obj = self._obj.clone().detach().numpy()
            return obj

        def obj_func_parallel(params):
            # # WRITE CODE HERE within two '#' bar 						#
            # # Converts the input params to list x form and dim(x)=2 	#
            # ###########################################################
            x = params
            # ###########################################################
            keylib = []
            reslib = []
            newlib = []
            f = []
            CV = []
            params = list(get_circuit().variables.keys())
            get_circuit().datasweep = DataSweep("sweepdata", params, x.T)
            with get_circuit().setenv(shell=False):
                res = self.tester.build()
                for k, v in res.items():
                    keylib.append(k)
                    reslib.append(v)
                reslib = Tensor(reslib).T
                for j in range(len(reslib)):
                    newdict = {}
                    for index, value in enumerate(keylib):
                        newdict[value] = TorchVariable(reslib[j][index])
                    newlib.append(newdict)
            test = [self.problem(z) for z in newlib]
            f_ret = [x.objective[0] for x in test]
            cv = [[list(x.constraint)] for x in test]
            CV = np.vstack(cv)
            f = np.array(f_ret)
            penaly_ueq = np.maximum(CV, 0).sum(axis=1)
            return f, penaly_ueq

        def tracer_updata(x):
            if self.tracer != None:
                tester = self.update(x)
                prob = self.problem(tester)
                self.tracer.collect_n_iters_fom(
                    prob.objective[0].clone().detach().numpy(),
                    np.array([con.to_numpy() for con in prob.constraint]),
                )

        def constrains(x):
            CV = []
            for i in range(len(x)):
                tester = self.update(x[i])
                prob = self.problem(tester)
                cv = [x.tolist() for x in prob.constraint]
                CV.append(cv)
                if self.tracer != None:
                    self.tracer.collect_n_iters_fom(
                        prob.objective[0].clone().detach().numpy(),
                        np.array([con.to_numpy() for con in prob.constraint]),
                    )

            CV = np.vstack(CV)
            penaly_ueq = np.maximum(CV, 0).sum(axis=1)
            return penaly_ueq

        self.method = LSHADE(
            fitness=obj_func_parallel,
            lowwer=lb,
            upper=ub,
            pop_size=20,
            dim=Dim,
            epochs=10000,
            tracer=self.tracer,
            eval=eval
        )

    def solve(self, prophetVars=None):
        '''
        prophetVars:
        '''
        best_x, tracer_t = self.method.run()
        iters_info = (
            (tracer_t.n_iters, tracer_t.n_iters_fom)
            if self.tracer != None
            else (None, None)
        )
        # iters_info = ((None, None))

        return {
            "success": True,
            "Vars": [best_x],
            "fom": iters_info,
        }

        pass

    def update(self, x):
        Variable.update_all_variables(x)
        with get_circuit().setenv(quite=True, shell=False):
            return self.tester.build()
